
from pymongo import UpdateOne
import tushare as ts
from pev3.util.DataBase import DB_CONN
from pev3.util.StockUtil import getTradingDateStockCodes
import urllib3
import json

"""
从东财获取财务数据，保存到数据库中
"""
class FinanceReportCrawler:

    def __init__(self):
        self.db = DB_CONN

    """
    从东方财富网站抓取三张财务报表
    :return:
    """
    def crawlFinanceReport(self):
        # 先获取所有的股票列表
        codes = getTradingDateStockCodes()

        # 创建连接池
        connPool = urllib3.PoolManager()

        # 抓取的网址，两个替换参数 {1} - 财报类型 {2} - 股票代码
        url = 'http://dcfm.eastmoney.com/em_mutisvcexpandinterface/api/js/get?' \
              'type={1}&token=70f12f2f4f091e459a279469fe49eca5&' \
              'filter=(scode={2})&st=reportdate&sr=-1&p=1&ps=500&' \
              'js={%22pages%22:(tp),%22data%22:%20(x)}&rt=51327674'

        userAgent = 'Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/70.0.3538.67 Safari/537.36'

        # 对应的类型，分别为资产负债表、现金流量表和利润表
        reportTypes = ['CWBB_ZCFZB', 'CWBB_XJLLB', 'CWBB_LRB']
        count = 0

        for code in codes:

            for reportType in reportTypes:
                print('开始抓取财报数据, 股票：%s, 财报类型：%s' % (code, reportType), flush=True)
                response = connPool.request('GET',
                                            url.replace('{1}', reportType).replace('{2}', code),
                                            headers={
                                                'User-Agent': userAgent
                                            })

                result = response.data
                # 解析抓取结果
                result = json.loads(response.data.decode('UTF-8'))
                reports = result['data']

                updateRequests = []
                for report in reports:
                    # 更新字段
                    try:
                        report.update({
                            # 公告日和报告期保留年月日
                            'announced_date': report['noticedate'][0:10],
                            'report_date': report['reportdate'][0:10],
                            # 股票名称和股票代码的字段名和系统设计保持一致
                            'code': code,
                            'name': report['sname']
                        })

                        updateRequests.append(
                            UpdateOne({
                                'code': code,
                                'report_date': report['report_date'],
                                'announced_date': report['announced_date']
                                },
                                {'$set': report},
                                upsert=True
                            )
                        )
                    except:
                        print('解析出错，股票: %s，财报类型: %s' % (code, reportType))

                if len(updateRequests) > 0:
                    updateResult = self.db[reportType].bulk_write(updateRequests, ordered=False)
                    print("股票：%s, 财报类型：%s, 插入：%4d, 更新: %4d" % (code, reportType, updateResult.upserted_count, updateResult.modified_count))

    """
    获取回测周期内的年度财报数据，组合成一个dict数据结构，key是股票代码，
    value是一个按照报告发布日期排序的列表，列表内的元素也是一个dict
    {'eps': 每股收益, 'announced_date': 公告日期}
    """
    def getCodeReports(self):

        # 这个tuple包含了三个元素，前两个分别是用来获取年度财报时的参数年份和季度，
        # 后面是一个发布年度财报的年份，因为一般情况下发布财报都是在第二年的4月底之前，
        # 所以这个年份比财报的年份晚一年
        reportDateTuples = [(2015, 4, 2016),(2016, 4, 2017),(2017, 4, 2018)]

        # 要返回的数据结构
        codeReportDict = dict()

        # 循环获取所有指定报告期的数据
        for reportDateTuple in reportDateTuples:
            # 从Tushare获取年报数据
            dfReports = ts.get_report_data(reportDateTuple[0], reportDateTuple[1])

            # 只需要股票代码、每股收益和公告日期三个字段
            codes = dfReports['code']
            epses = dfReports['eps']
            announced_dates = dfReports['report_date']

            # 这个是报告发布的年度
            announcedYear = str(reportDateTuple[2])
            

            # 拿到已经缓存的股票代码集合
            codesOfCachedReports = set(codeReportDict.keys())

            for index in dfReports.index:
                code = codes[index]
                eps = epses[index]
                announcedDate = announced_dates[index]

                print('%s %5.2f %d %s' % (code, eps, reportDateTuple[0], announcedDate), flush=True)

                # 如果eps是非数字，或者发布日期的月份超过了4月，就不作处理，因为股票在上市前
                # 也会发布财报，那么这个财报的发布日期可能不是定期报告所规定的时间范围，
                # 那么对这种上市之前的数据暂时不予处理
                if str(eps) != 'nan' and int(announcedDate[0:2]) <= 4:
                    # 组合成完整的公告年月日
                    announcedDate = announcedYear + '-' + announcedDate
                    print('%s %5.2f %s' %(code, eps, announcedDate), flush=True)

                    # 如果当前股票不在需要返回的数据结构中，则添加到其中
                    if code not in codesOfCachedReports:
                        codeReportDict[code] = []
                        codesOfCachedReports.add(code)

                    # 将eps和公告日期添加到列表中
                    codeReportDict[code].append({'eps': eps, 'announced_date': announcedDate})
        # 返回获取的数据
        return codeReportDict

if __name__ == '__main__':
    # FinanceReportCrawler().crawlFinanceReport()
    print(FinanceReportCrawler().getCodeReports())